package stateful;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.api.java.tuple.Tuple4;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction;
import org.apache.flink.util.Collector;

/**
 * 广播状态 动态的广播变量
 */
public class BroadcastStateDemo {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());

        env.enableCheckpointing(3000);

        // 维度流
        DataStreamSource<String> dimStream = env.socketTextStream("hadoop1", 8888);
        // 事实流
        DataStreamSource<String> truthStream = env.socketTextStream("hadoop1", 9999);

        SingleOutputStreamOperator<Tuple3<String, String, String>> dimTupleStream = dimStream.map(new MapFunction<String, Tuple3<String, String, String>>() {
            @Override
            public Tuple3<String, String, String> map(String value) throws Exception {
                String[] fields = value.split(",");
                return Tuple3.of(fields[0], fields[1], fields[2]);
            }
        });

        // 创建状态描述对象
        MapStateDescriptor<String, String> broadcastMapStateDescriptor = new MapStateDescriptor<>("dim-state", String.class, String.class);

        // 维度流转换为广播流
        BroadcastStream<Tuple3<String, String, String>> broadcastStream = dimTupleStream.broadcast(broadcastMapStateDescriptor);

        // 处理事实流
        SingleOutputStreamOperator<Tuple3<String, String, Double>> truthTpStream = truthStream.map(new MapFunction<String, Tuple3<String, String, Double>>() {
            @Override
            public Tuple3<String, String, Double> map(String value) throws Exception {
                String[] fields = value.split(",");
                return Tuple3.of(fields[0], fields[1], Double.parseDouble(fields[2]));
            }
        });

        KeyedStream<Tuple3<String, String, Double>, Tuple2<String, String>> keyedStream = truthTpStream.keyBy(t -> Tuple2.of(t.f0, t.f1), TypeInformation.of(new TypeHint<Tuple2<String, String>>() {
        }));

        // 连接
        keyedStream.connect(broadcastStream)
                .process(new MyProcessFunction(broadcastMapStateDescriptor))
                .print();

        env.execute("");
    }

    private static class MyProcessFunction extends KeyedBroadcastProcessFunction<Tuple2<String, String>, Tuple3<String, String, Double>, Tuple3<String, String, String>, Tuple4<String, String, String, Double>> {

        // 维度状态描述信息
        private MapStateDescriptor<String, String> broadcastMapStateDescriptor;


        public MyProcessFunction(MapStateDescriptor<String, String> broadcastMapStateDescriptor) {
            this.broadcastMapStateDescriptor = broadcastMapStateDescriptor;
        }

        /**
         * 处理事实数据 与维度数据关联 参与活动的人数
         *
         * @param t
         * @param ctx
         * @param out
         * @throws Exception
         */
        @Override
        public void processElement(Tuple3<String, String, Double> t, ReadOnlyContext ctx, Collector<Tuple4<String, String, String, Double>> out) throws Exception {
            // 获取广播状态
            ReadOnlyBroadcastState<String, String> broadcastState = ctx.getBroadcastState(broadcastMapStateDescriptor);
            String uid = t.f0;
            String eventId = t.f1;
            Double money = t.f2;
            // 获取用户信息
            String eventName = broadcastState.get(eventId);
            out.collect(Tuple4.of(uid, eventId, eventName, money));
        }


        /**
         * 处理维度数据
         *
         * @param t
         * @param ctx
         * @param out
         * @throws Exception
         */
        @Override
        public void processBroadcastElement(Tuple3<String, String, String> t, Context ctx, Collector<Tuple4<String, String, String, Double>> out) throws Exception {
            // INSERT UPDATE DELETE
            String type = t.f0;
            String eventId = t.f1;
            String eventName = t.f2;
            // 获取广播状态
            BroadcastState<String, String> broadcastState = ctx.getBroadcastState(broadcastMapStateDescriptor);
            if ("DELETE".equals(type)) {
                broadcastState.remove(eventId);
            } else {
                broadcastState.put(eventId, eventName);
            }
        }
    }
}
